from typing import List, Optional, Tuple, Sequence
import torch


def _to_mask_matrix(aw: Sequence, W: Optional[int] = None) -> List[List[int]]:
    """Convert aw input (ints or binary rows) into List[List[int]] mask matrix.
    aw can be:
      - torch.Tensor or list of ints (each int is a bitmask)
      - torch.Tensor or list of lists of 0/1 (already bit rows)
    W: optional bit width; if None it is inferred from max int's bit_length or len(rows)
    """
    # If tensor
    if isinstance(aw, torch.Tensor):
        a = aw.detach().cpu()
        if a.dim() == 1 and a.dtype in (torch.long, torch.int):
            ints = [int(x.item()) for x in a]
            if W is None:
                maxv = max(ints) if len(ints) > 0 else 1
                W = max(1, maxv.bit_length())
            return [ [(val >> j) & 1 for j in range(W)] for val in ints ]
        else:
            # treat as 2D mask
            m = a.numpy().tolist()
            # ensure ints 0/1
            rows = [[int(x) for x in row] for row in m]
            if W is None:
                W = len(rows[0]) if rows else 0
            return [row[:W] + [0]*(W - len(row)) if len(row) < W else row for row in rows]

    # If list-like
    try:
        # check element types
        first = aw[0]
    except Exception:
        return []
    # list of ints
    if isinstance(first, int):
        ints = [int(x) for x in aw]
        if W is None:
            maxv = max(ints) if len(ints) > 0 else 1
            W = max(1, maxv.bit_length())
        return [ [(val >> j) & 1 for j in range(W)] for val in ints ]
    # list of sequences -> binary rows
    rows = [ [int(x) for x in row] for row in aw ]
    if W is None:
        W = len(rows[0]) if rows else 0
    return [row[:W] + [0]*(W - len(row)) if len(row) < W else row for row in rows]


def schedule_from_aw(aw: Sequence, W: Optional[int] = None) -> Optional[List[int]]:
    """Given per-node available-window masks, return an assignment of nodes -> distinct time slots.

    Input:
      - aw: sequence (list or 1D tensor) of integers (bitmasks) or sequence of binary rows.
      - W: optional bit width to interpret ints or to pad/truncate rows.

    Output:
      - list assignment of length k: assignment[i] = slot index (0..W-1) assigned to node i
        if assignment exists; otherwise returns None when no full matching exists.

    The algorithm uses bipartite matching (Kuhn/DFS) to assign each node to a distinct time slot
    where its mask has a 1. If sum(OR bits) < k it immediately returns None.
    """
    mat = _to_mask_matrix(aw, W)
    k = len(mat)
    if k == 0:
        return []
    W = len(mat[0])

    # quick necessary condition
    or_bits = [0] * W
    for row in mat:
        for j, b in enumerate(row):
            or_bits[j] |= (1 if b else 0)
    if sum(or_bits) < k:
        return None

    # match_slot[s] = node index assigned to slot s, or -1
    match_slot: List[int] = [-1] * W

    def try_assign(u: int, seen: List[bool]) -> bool:
        for s, allowed in enumerate(mat[u]):
            if not allowed or seen[s]:
                continue
            seen[s] = True
            if match_slot[s] == -1 or try_assign(match_slot[s], seen):
                match_slot[s] = u
                return True
        return False

    for u in range(k):
        seen = [False] * W
        if not try_assign(u, seen):
            return None

    # invert match_slot to node->slot mapping
    assignment = [-1] * k
    for s, u in enumerate(match_slot):
        if u != -1:
            assignment[u] = s
    return assignment


def all_schedules_from_aw(aw: Sequence, W: Optional[int] = None, max_solutions: int = 1000) -> Optional[List[List[int]]]:
    """Return all possible full assignments (node -> slot) respecting aw masks.

    - aw: sequence of ints (bitmasks) or binary rows.
    - W: optional width (number of time slots). If None inferred.
    - max_solutions: cap on number of solutions to enumerate (to avoid combinatorial blowup).

    Returns a list of assignments (each a list of length k with slot indices) or
    an empty list if none exist. If enumeration was cut short due to max_solutions,
    returns the found subset.
    """
    mat = _to_mask_matrix(aw, W)
    k = len(mat)
    if k == 0:
        return []
    W = len(mat[0])

    # quick necessary condition
    or_bits = [0] * W
    for row in mat:
        for j, b in enumerate(row):
            or_bits[j] |= (1 if b else 0)
    if sum(or_bits) < k:
        return []

    # order nodes by increasing availability (pruning)
    nodes_order = sorted(range(k), key=lambda u: sum(mat[u]))
    assigned = [-1] * k
    used = [False] * W
    solutions: List[List[int]] = []

    def backtrack(pos: int):
        if len(solutions) >= max_solutions:
            return
        if pos == k:
            solutions.append(assigned.copy())
            return
        u = nodes_order[pos]
        # iterate available slots for node u
        for s in range(W):
            if mat[u][s] and not used[s]:
                used[s] = True
                assigned[u] = s
                backtrack(pos + 1)
                if len(solutions) >= max_solutions:
                    return
                used[s] = False
                assigned[u] = -1

    backtrack(0)
    return solutions


def is_unique_schedule(aw: Sequence, W: Optional[int] = None) -> Tuple[bool, Optional[List[int]]]:
    """Return (is_unique, assignment_or_none).

    - is_unique True if exactly one full assignment exists; False otherwise.
    - assignment_or_none: the unique assignment if is_unique True, else one found assignment or None.
    """
    sols = all_schedules_from_aw(aw, W=W, max_solutions=2)
    if not sols:
        return False, None
    if len(sols) == 1:
        return True, sols[0]
    return False, sols[0]


# quick smoke tests when run as script
if __name__ == '__main__':
    # Example 1: simple schedulable
    aw1 = [0b0011, 0b0101, 0b1000]
    asg = schedule_from_aw(aw1)
    print('aw1 assignment:', asg)

    # Example 2: unschedulable (only one slot available)
    aw2 = [0b001, 0b001, 0b001]
    print('aw2 assignment (expect None):', schedule_from_aw(aw2))

    # Example 3: rows provided
    rows = [[0,0,0,1],[0,0,1,0],[1,0,0,0]]
    print('rows assignment:', schedule_from_aw(rows))
    print('rows assignment:', all_schedules_from_aw(rows))
    # new enumeration example (example3 like multiple solutions)
    # Suppose masks: three nodes with overlapping availability
    aw3 = [0b0111, 0b1011, 0b1101]  # these may admit multiple matchings
    all_sols = all_schedules_from_aw(aw3, max_solutions=100)
    print('aw3 all solutions count:', len(all_sols))
    for sol in all_sols[:10]:
        print('  sol:', sol)
    unique, one = is_unique_schedule(aw3)
    print('aw3 unique:', unique, 'example:', one)
